#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
用于读取traceSeq生成的h5文件并绘制相关图表的脚本
"""

import argparse
import os,sys
import numpy as np
from pylab import plt
import pandas as pd
from obspy.core.utcdatetime import UTCDateTime

# 导入自定义工具
from utils.h5data import read_h5_data
from utils.trace import filtfilt
from utils.plot import plot_traces, plot_raw_data
from utils.loc import load_loc, get_distance
from utils.math import norm

# 命令行参数解析
parser = argparse.ArgumentParser(description='Plot traceSeq h5 data')
parser.add_argument('-input', default='', help='input h5 file')
parser.add_argument('-figroot', default='figures/6.traceSeq.figures', help='root to save figs')
args = parser.parse_args()

INPUT_FILE = args.input
FIG_ROOT = args.figroot

datasets,args_infile = read_h5_data(INPUT_FILE, keys=['MARKER','EMARKER'], group_name='metadata',read_attrs=True)
print(args_infile)
DATE=args_infile['date']
fs = args_infile['fs']
fe = args_infile['fe']

EMARKER = datasets[1].decode()
MARKER = datasets[0].decode()
FIG_ROOT = f'{FIG_ROOT}/{DATE}.{EMARKER}.{MARKER}'
if not os.path.exists(FIG_ROOT):
    os.mkdir(FIG_ROOT)

# 读取h5文件中的数据
print(f'Reading data from {INPUT_FILE}')
data_dict = read_h5_data(INPUT_FILE, keys=['t', 'ref', 'S', 'R','x'])
tn, ref_traces,S, R,x = data_dict
# 解码字符串数组
CORR_STATS = [name.decode() for name in R]
S = S.decode()

# 绘制参考道数据
print('Plotting reference traces...')

ns = len(CORR_STATS)

# 绘制参考道数据和频散图
# if args_infile['Tpos']=='N':
#     x=-x

fs,fe=3,15
VLIM = [-5000, 5000]
TLIM=[-5,5]

# fs,fe=1,2.5
# VLIM = [100, 1000]
# TLIM=[-7,7]
# v_estimate = 300

# ref_traces[np.where(x>-500)[0],:]=0

NORM=False
# PLOT_WIGGLE = True if ns<30 else False
PLOT_WIGGLE = True
print(ns)
ref_traces = filtfilt(ref_traces,tn[1]-tn[0],[fs,fe], order=4)
ref_traces = norm(ref_traces, ONE_AXIS=True)
fig, ax1, ax2 = plot_raw_data(ref_traces, x, tn, fs=fs, fe=fe, VLIM=VLIM, PLOT_WIGGLE=PLOT_WIGGLE, SCALE=60, FV_NORM=NORM)
if PLOT_WIGGLE:
    for j in range(ns):
        ax1.text(-1, x[j], CORR_STATS[j], fontsize=6)
    ax1.set_ylim([x.min(),x.max()])
else:
    ax1.set_ylim([x.min(),x.max()])
    plt.clim()
    

ax1.plot([-5, 5], [-5*3000, 5*3000])
ax1.plot([-5, 5], [-5*340, 5*340])
ax1.set_title(f'{S}.{MARKER}')
ax1.set_xlim(TLIM)

# 频散曲线
disp = np.loadtxt('figures/6.traceSeq.figures/disp.txt',
                  delimiter=',')
ax2.plot(disp[:,0],disp[:,1],'k',lw=1)

# f0=300/3.6/25
# for j in range(int(fe/f0)+1):
#     ax2.plot([j*f0,j*f0],VLIM,color='k',lw=1)

# tau =32/80+32/3000
# f0=1/tau
# for j in range(int(fe/f0)+1):
#     ax2.plot([j*f0,j*f0],VLIM,color='k',linestyle='--',lw=1)
f0=4.4
for j in [1,4,9]:
    ax2.plot([j*f0,j*f0],VLIM,color='yellow',lw=1)
ax2.set_ylim(VLIM)
fig.tight_layout()
figname = f'{FIG_ROOT}/ref.{fs:03.1f}.{fe:03.1f}.NORM{int(NORM)}.png'
fig.savefig(figname, dpi=300)
print(f'Saved {figname}')

print('Plotting completed.')